'''
StreamEvent 事件流处理示例

StreamEvent 是 LangChain 提供的高级特性，允许你实时监控大模型生成过程的每个阶段。

安装依赖：
pip install "langserve[all]"  # 这会安装 langserve, langchain, 以及所有常用依赖

关键概念解析：
1. Pydantic：Python 中用于数据验证和解析的第三方库
2. LangServe：快速部署 LangChain 项目为 API 的服务框架
3. StreamEvent：大模型生成过程中的实时事件流

LangServe 核心功能：
    - 快速部署：几行代码就能将 LangChain 项目变成 API
    - 自动生成文档：自动生成 OpenAPI 文档，让 API 使用者一目了然
    - 内置客户端：提供 LangChain.js 客户端，方便前端 JavaScript 调用
    - 生产就绪：基于 FastAPI，具备高性能、异步支持等生产级特性
    - Playground：为每个部署的链自动提供一个简单的 Web UI，方便测试和调试

LangChain CLI 使用：
    pip install -U langchain-cli
    安装后可以执行 langchain 命令查看所有可用命令
    示例：langchain app new myserve  # 快速创建新的 LangServe 项目
'''

import asyncio
import os

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI

# 自动加载 .env 文件中的环境变量
load_dotenv()

# 定义支持流式输出的模型
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="deepseek-ai/DeepSeek-R1-0528", streaming=True, temperature=0.7)


async def main():
    """
    演示 StreamEvent 事件流的异步处理

    astream_events 方法会返回一个异步生成器，实时产生大模型生成过程中的各种事件
    每个事件都包含当前处理阶段的状态信息
    """
    events = []  # 存储所有事件的列表，用于后续分析

    print("=== 开始监听大模型生成事件流 ===")
    print("事件类型说明:")
    print("- on_chat_model_start: 聊天模型开始处理")
    print("- on_chat_model_stream: 聊天模型流式输出内容")
    print("- on_chat_model_end: 聊天模型处理结束")
    print("- 其他事件: 各种中间处理步骤")
    print("=" * 50)

    # 核心：监听大模型生成过程中的所有事件
    async for event in llm.astream_events("帮我写一首关于牛的短诗", version="v2"):
        """
        astream_events 参数说明：
        - 第一个参数：输入给模型的提示词
        - version="v2": 使用第二版事件格式（更稳定和功能完整）
        
        事件对象通常包含：
        - event: 事件类型（如 'on_chat_model_stream'）
        - name: 组件名称
        - run_id: 运行ID（用于追踪）
        - tags: 标签
        - metadata: 元数据
        - data: 事件具体数据（如生成的文本内容）
        """
        print(f"\n🔔 事件类型: {event['event']}")
        print(f"   运行ID: {event['run_id']}")
        print(f"   组件: {event['name']}")

        # 根据不同事件类型显示特定信息
        if event['event'] == 'on_chat_model_stream':
            # 流式输出事件，包含实际生成的内容
            if 'chunk' in event['data'] and hasattr(event['data']['chunk'], 'content'):
                content = event['data']['chunk'].content
                print(f"   📝 生成内容: '{content}'")

        elif event['event'] == 'on_chat_model_start':
            # 模型开始处理事件
            print(f"   🚀 开始处理请求")
            if 'messages' in event['data']:
                print(f"   输入消息: {event['data']['messages']}")

        elif event['event'] == 'on_chat_model_end':
            # 模型处理结束事件
            print(f"   ✅ 处理完成")
            if 'output' in event['data']:
                print(f"   最终输出: {event['data']['output']}")

        # 存储事件用于后续分析
        events.append(event)

    print("\n" + "=" * 50)
    print("=== 事件流分析总结 ===")
    print(f"总共捕获事件数: {len(events)}")

    # 统计不同类型事件的数量
    event_types = {}
    for event in events:
        event_type = event['event']
        event_types[event_type] = event_types.get(event_type, 0) + 1

    print("事件类型分布:")
    for event_type, count in event_types.items():
        print(f"  {event_type}: {count}次")

    # 提取所有生成的内容
    generated_content = ""
    for event in events:
        if (event['event'] == 'on_chat_model_stream' and
                'chunk' in event['data'] and
                hasattr(event['data']['chunk'], 'content')):
            generated_content += event['data']['chunk'].content

    if generated_content:
        print(f"\n📖 完整生成内容: {generated_content}")


# 运行异步主函数
if __name__ == "__main__":
    asyncio.run(main())
